#!/usr/bin/env python3
"""
OpenAI API integration for generating Manim code.

This module provides functionality to generate Manim animation code using the OpenAI API.
It handles API communication, response parsing, and code extraction.

Environment Variables:
    OPENAI_API_KEY: Your OpenAI API key (required)

Usage:
    from openai_manim_generator import generate_manim_code, save_code_to_file

    code = generate_manim_code("Create a circle animation")
    save_code_to_file(code, "my_animation.py")
"""

import os
import sys
import argparse
import requests
from pathlib import Path
from typing import Optional


from dotenv import load_dotenv
import os
load_dotenv()
import openai

API_BASE_URL=os.getenv("OPENAI_BASE_URL", "")
DEFAULT_MODEL=os.getenv("MODEL", "")
TEMPERATURE = 0.2

def get_api_key() -> Optional[str]:
    """
    Get the OpenAI API key from environment variables.

    Returns:
        The API key if found, None otherwise
    """
    return os.environ.get("OPENAI_API_KEY", "YOUR_OPEN_API_KEY")


def create_system_prompt() -> str:
    """Create the system prompt for the OpenAI API."""
    return """You are an expert Python programmer specializing in the Manim library for creating mathematical animations. Generate complete, working Manim code that creates clear and educational animations."""


def create_user_message(prompt: str) -> str:
    """
    Create the user message for the OpenAI API.

    Args:
        prompt: The user's animation request

    Returns:
        Formatted user message
    """
    base_message = f"""I need Python code using the Manim library (version 0.17.3) to create an animation that explains the following concept:

{prompt}"""

    base_message += """

Please generate complete, working Manim code that creates a clear and educational animation. The code should:

1. Import the necessary Manim modules
2. You can use matplotlib, numpy, scipy, sympy, and networkx if needed
3. Define a Scene class that extends Scene
4. Include a construct method that builds the animation
5. Use appropriate Manim objects and animations to visualize the concept
6. Include helpful comments explaining the code
7. Do not use LaTeX or any text that requires LaTeX rendering. Use Manim's Text class instead of Tex or MathTex.
8. Be complete and ready to run without additional modifications
9. IMPORTANT: Use Create() instead of ShowCreation() as ShowCreation is deprecated in Manim 0.17.3
10. IMPORTANT: Use Write() for text animations and FadeIn() for other objects appearing
11. IMPORTANT: For any animations, check that they exist in Manim 0.17.3
12. The code should be well-organized, easy to read, and follow best practices for Python and Manim.
13. The code should be surrouded by de <code> </code> tags.

Only provide the Python code without any additional explanations or text outside the code blocks."""

    return base_message


def make_api_request(prompt: str, model: str) -> Optional[dict]:
    """
    Make a request to the OpenAI API using the official Python client.

    Args:
        prompt: The animation prompt
        model: The model to use

    Returns:
        API response data or None if failed
    """
    import os
    from openai import OpenAI
    from openai import APIConnectionError, APITimeoutError, RateLimitError, APIError
    
    api_key = get_api_key()
    if not api_key:
        print("Error: OPENAI_API_KEY environment variable not set")
        return None
    
    base_url = os.environ.get("OPENAI_BASE_URL", "https://api.openai.com/v1")
    
    try:
        # Create OpenAI client instance
        client = OpenAI(
            api_key=api_key,
            base_url=base_url.rstrip('/')  # Ensure no trailing slash
        )
        
        system_prompt = create_system_prompt()
        user_message = create_user_message(prompt)
        
        print(f"\n=== Sending prompt to API ===\n"
              f"Base URL: {base_url}\n"
              f"Model: {model}\n"
              f"System: {system_prompt[:80]}...\n"
              f"User: {user_message[:200]}...\n==================")
        
        # Make API request
        response = client.chat.completions.create(
            model=model,
            messages=[
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": user_message}
            ],
            temperature=TEMPERATURE
        )
        
        # Format response into dictionary for backward compatibility
        return {
            "id": response.id,
            "object": response.object,
            "created": response.created,
            "model": response.model,
            "choices": [
                {
                    "index": choice.index,
                    "message": {
                        "role": choice.message.role,
                        "content": choice.message.content
                    },
                    "finish_reason": choice.finish_reason
                }
                for choice in response.choices
            ],
            "usage": dict(response.usage),
        }
    
    except APIConnectionError as e:
        print(f"Network connection error: {e}")
        return None
    except APITimeoutError as e:
        print(f"API timeout error: {e}")
        return None
    except RateLimitError as e:
        print(f"Rate limit exceeded: {e}")
        return None
    except APIError as e:
        print(f"API response error: HTTP {e.status_code} - {e.message}")
        return None
    except Exception as e:
        print(f"Unexpected error during API request: {e}")
        return None


def extract_code_from_response(response_data: dict) -> str:
    """
    Extract code from the OpenAI API response.

    Args:
        response_data: The API response data

    Returns:
        Extracted code or empty string if failed
    """
    if 'choices' not in response_data or len(response_data['choices']) == 0:
        print(f"Unexpected response format: {response_data}")
        return ""

    # Get content from the first choice's message
    generated_content = response_data['choices'][0]['message']['content']
    print(generated_content)
    # Clean up the response to extract just the code
    if "<code>" in generated_content:
        code_parts = generated_content.split("<code>")
        if len(code_parts) > 1:
            # Take the last code block in case of multiple entries
            code_block = code_parts[-1].split("</code>")[0]
            generated_content = code_block.strip()
    
    if "```python" in generated_content:
        code_parts = generated_content.split("```python")
        if len(code_parts) > 1:
            # Take the last code block in case of multiple entries
            code_block = code_parts[-1].split("```")[0]
            return code_block.strip()
    
    # Fallback to everything between first ``` and last ```
    if "```" in generated_content:
        code_start = generated_content.find("```") + 3
        code_end = generated_content.rfind("```")
        if code_end > code_start:
            return generated_content[code_start:code_end].strip()
    print("=====================================================")
    print(generated_content)
    # Fallback to entire content
    return generated_content.strip()


def generate_manim_code(prompt: str, model: str = DEFAULT_MODEL) -> str:
    """
    Generate Manim code using the OpenAI API.

    Args:
        prompt: The prompt describing the animation to generate
        model: The OpenAI model to use

    Returns:
        The generated Manim code as a string, or empty string if failed
    """
    try:
        response_data = make_api_request(prompt, model)
        if not response_data:
            return ""

        return extract_code_from_response(response_data)

    except Exception as e:
        print(f"Error generating code: {e}")
        return ""


def save_code_to_file(code: str, output_file: str = "generated_animation.py") -> str:
    """
    Save the generated code to a file.

    Args:
        code: The code to save
        output_file: The file to save the code to

    Returns:
        The path to the saved file, or empty string if failed
    """
    if not code.strip():
        print("Error: No code to save")
        return ""

    try:
        output_path = Path(output_file)
        output_path.parent.mkdir(parents=True, exist_ok=True)

        with open(output_path, "w", encoding="utf-8") as f:
            f.write(code)

        print(f"Code successfully saved to {output_path}")
        return str(output_path)

    except Exception as e:
        print(f"Error saving code to file: {e}")
        return ""


def main():
    """Command-line interface for the module."""
    parser = argparse.ArgumentParser(description="Generate Manim code using OpenAI API")
    parser.add_argument("prompt", type=str, help="Prompt describing the animation to generate")
    parser.add_argument("--output", "-o", type=str, default="generated_animation.py",
                        help="Output file for the generated code")
    parser.add_argument("--model", "-m", type=str, default=DEFAULT_MODEL,
                        help="OpenAI model to use")

    args = parser.parse_args()

    print(f"Generating Manim code for: {args.prompt}")
    generated_code = generate_manim_code(args.prompt, args.model)

    if not generated_code:
        print("Failed to generate code.")
        return 1

    output_file = save_code_to_file(generated_code, args.output)
    if not output_file:
        print("Failed to save code.")
        return 1

    print(f"Code saved to {output_file}")
    return 0


if __name__ == "__main__":
    sys.exit(main())
